package main

import (
	"errors"
	"flag"
	"fmt"
	"wisdom_education_in_kunshan/znet"
	"log"
	"math/rand"
	"net/http"
	"net/http/cookiejar"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// Task ...
type Task struct {
	url       string
	rArg      bool
	rArgName  string
	dNum      int
	ddPath    []string
	force     bool
	threadNum int
	plugin    bool
}

var task Task
var threadNum chan *http.Client
var wg sync.WaitGroup

func main() {
	_url := flag.String("url", "", "验证码下载地址 必须")
	_randURLArg := flag.Bool("r", false, "是否在url之后附加随机值\n默认值:false")
	_randURLArgName := flag.String("a", "", "url之后附加随机值：https://test.com?${rname}=232142\n如果为空:https://test.com?232142\n默认值:空")
	_downloadNum := flag.Int("n", 1, "验证码下载个数\n默认值:1")
	_downloadPath := flag.String("o", "./tmp", "验证码下载存储路径\n默认值:./tmp")
	_force := flag.Bool("f", false, "强制执行,会先清空下载目录")
	_threadNum := flag.Int("t", 3, "下载最大线程数(1~5)\n默认值:3")
	_plugin := flag.Bool("p", false, "插件")
	flag.Parse()
	task.url = strings.Trim(*_url, " ")
	if task.url == "" {
		log.Fatalln("下载url不能为空")
	}
	task.force = *_force
	task.rArg = *_randURLArg
	task.rArgName = strings.Trim(*_randURLArgName, " ")
	task.dNum = *_downloadNum
	if task.dNum <= 0 {
		task.dNum = 1
	}
	task.threadNum = *_threadNum
	if task.threadNum > task.dNum {
		task.threadNum = task.dNum
	}
	if task.threadNum < 0 {
		task.threadNum = 1
	}
	if task.threadNum > 5 {
		task.threadNum = 5
	}
	folder := strings.Trim(*_downloadPath, " ")
	if folder == "" {
		folder = "./"
	}
	var err error
	folder, err = filepath.Abs(folder)
	if err != nil {
		log.Fatalln("检测下载目录出错", err)
	}
	task.plugin = *_plugin
	if task.plugin {
		task.ddPath = []string{folder + string(os.PathSeparator) + "src", folder + string(os.PathSeparator) + "dst"}
	} else {
		task.ddPath = []string{folder}
	}
	err = checkDirs()
	if err != nil {
		log.Fatalln("检测下载目录出错", err)
	}
	threadNum = make(chan *http.Client, task.threadNum)
	for i := 0; i < task.threadNum; i++ {
		threadNum <- generateClient()
	}
	log.Println("开始下载任务", task)
	for i := 0; i < task.dNum; i++ {
		wg.Add(1)
		client := <-threadNum
		go downcode(client)
	}
	wg.Wait()
}

func generateURL() string {
	if task.rArg {
		switch strings.LastIndex(task.url, "?") {
		case -1:
			if task.rArgName == "" {
				return fmt.Sprintf("%s?%d", task.url, rand.Intn(65535))
			}
			return fmt.Sprintf("%s?%s=%d", task.url, task.rArgName, rand.Intn(65535))
		case 0:
			if task.rArgName == "" {
				return fmt.Sprintf("%s%d", task.url, rand.Intn(65535))
			}
			return fmt.Sprintf("%s%s=%d", task.url, task.rArgName, rand.Intn(65535))
		default:
			switch strings.LastIndex(task.url, "&") {
			case 0:
				if task.rArgName == "" {
					return fmt.Sprintf("%s%d", task.url, rand.Intn(65535))
				}
				return fmt.Sprintf("%s%s=%d", task.url, task.rArgName, rand.Intn(65535))
			default:
				if task.rArgName == "" {
					return fmt.Sprintf("%s&%d", task.url, rand.Intn(65535))
				}
				return fmt.Sprintf("%s&%s=%d", task.url, task.rArgName, rand.Intn(65535))
			}
		}
	}
	return task.url
}

func generateClient() *http.Client {
	httpClient := znet.NewClient()
	jar, newHTTPErr := cookiejar.New(nil)
	if newHTTPErr != nil {
		log.Fatalln("Init httpClient err", newHTTPErr)
	}
	httpClient.Jar = jar
	httpClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		return errors.New("Not Allow Redirect")
	}
	return httpClient
}

func generateFileName() string {
	nt := time.Now()
	min := nt.Minute()
	sec := nt.Second()
	rid := (rand.Intn(10000) + 1000)
	return fmt.Sprintf("%02d%02d_%03d.png", min, sec, rid)
}
func checkDirs() (err error) {
	for _, dir := range task.ddPath {
		err = checkDir(dir)
		if err != nil {
			return
		}
	}
	return nil
}

func checkDir(dir string) error {
	var err error
	ok, err := PathExists(dir)
	if err != nil {
		return err
	}
	if ok && task.force {
		err = os.RemoveAll(dir)
		if err != nil {
			return err
		}
		ok = false
	}
	if !ok {
		ok, err = Mkdirs(dir)
		if err != nil {
			return err
		}
		if !ok {
			txt := fmt.Sprintf("创建目录(%s)出错", dir)
			return errors.New(txt)
		}
	}
	apt, err := ListDir(dir, FileOrDirFilter(true))
	if err != nil {
		return err
	}
	if !apt {
		txt := fmt.Sprintf("目录(%s)存在其他文件,如果指定此目录,请使用-f指令", dir)
		return errors.New(txt)
	}
	return nil
}

func downcode(httpClient *http.Client) {
	defer func() {
		threadNum <- httpClient
		wg.Done()
	}()
	// downCodeURL := fmt.Sprintf("http://sso.ksedu.cn/CodeImage.aspx?%d", rand.Intn(65535))
	downCodeURL := generateURL()
	log.Println(downCodeURL)
	req, _ := http.NewRequest("GET", downCodeURL, nil)
	req.Header.Add("Connection", "keep-alive")
	req.Header.Add("User-Agent", "Mozilla/5.0 (Windows NT 6.1; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36")
	res, err := httpClient.Do(req)
	if err != nil {
		log.Println(downCodeURL, err)
		return
	}
	defer res.Body.Close()
	fileName := generateFileName()
	srcFile := task.ddPath[0] + string(os.PathSeparator) + fileName
	err = SaveFile(srcFile, res.Body)
	if err != nil {
		log.Println("Save CodeImage err", err)
	}
	err = SavePNG(srcFile, srcFile)
	if err != nil {
		log.Println("Save CodeImage err", err)
	}
	if task.plugin {
		dstFile := task.ddPath[1] + string(os.PathSeparator) + fileName
		processCodeImage(srcFile, dstFile)
	}
}

func processCodeImage(path, savePath string) bool {
	err := Translate(path, savePath)
	if err != nil {
		log.Println("Translate Code Image err", err)
		return false
	}
	return true
}
